#!/usr/bin/python
# -*- coding: utf-8 -*-
import imp
import wx
from threading 		import *
from PortScanThread import PortScanThread
from NetStatThread  import NetStatThread
from LogsThread		import LogsThread
from Carry          import EVT_RESULT
import glob
import pygeoip
import matplotlib.pyplot as plt

db = imp.load_source('Database', 'database/db.py'); #uso de imp para abrir module na pasta database/


"""Frame Principal da GUI"""
class Frame(wx.Frame):
    def __init__(self):
        wx.Frame.__init__(self, None, wx.ID_ANY,title="WilFonScanner", pos=(200,200), size=(840,500),style=wx.DEFAULT_FRAME_STYLE ^ wx.RESIZE_BORDER)
        wx.Frame.CenterOnScreen(self)
        self.Bind(wx.EVT_CLOSE, self.OnClose)

        self.netstat    = None
        self.portscan   = None
        self.logs       = None
        self.db = db.Database()
        self.user = None
        self.gi = pygeoip.GeoIP('GeoIP.dat')

        iconFile = "lock.ico"
        icon1 = wx.Icon(iconFile, wx.BITMAP_TYPE_ICO)
        self.SetIcon(icon1)

        self.CreateStatusBar(style=0)
        self.SetStatusText("Criado por Miguel Fonseca & Wilson Pinto - MESI LPD 2013/2014")

        #---- PANEL -----#
        self.panel = wx.Panel(self,wx.ID_ANY)

        self.LoginBox()
        self.CreateBox()
        self.panel.Layout()
        #----- END OF PANEL -----#
        
        EVT_RESULT(self,self.OnResult)

    def hideLoginBox (self):
    	self.lbl_Username.Hide()
    	self.txt_Username.Hide()
    	self.lbl_Password.Hide()
    	self.txt_Password.Hide()
    	self.btn_Process.Hide()
    	self.l_login.Hide()
    	self.lbl_label.Hide()

    def hideCreateBox (self):
    	self.lbl2_Username.Hide()
    	self.txt2_Username.Hide()
    	self.lbl2_Password.Hide()
    	self.txt2_Password.Hide()
    	self.btn2_Process.Hide()
    	self.l2_login.Hide()
    	self.lbl2_label.Hide()
        self.pic.Hide()

    def LoginBox (self):
    	sizer = wx.FlexGridSizer(rows=3, cols=2, hgap=5, vgap=15)
    	self.l_login = wx.StaticText(self.panel, -1, "LOGIN",pos=(50,12))
    	self.l_login.SetFont(wx.Font(16, wx.SWISS, wx.NORMAL, wx.BOLD))

    	self.txt_Username = wx.TextCtrl(self.panel, 1, size=(150, -1))
    	self.lbl_Username = wx.StaticText(self.panel, -1, "Utilizador:")
    	sizer.Add(self.lbl_Username,0, wx.LEFT|wx.TOP| wx.RIGHT, 50)
    	sizer.Add(self.txt_Username,0, wx.TOP| wx.RIGHT, 50)

    	self.txt_Password = wx.TextCtrl(self.panel, 1, size=(150, -1), style=wx.TE_PASSWORD)
    	self.lbl_Password = wx.StaticText(self.panel, -1, "Password:")
    	sizer.Add(self.lbl_Password,0, wx.LEFT|wx.RIGHT, 50)
    	sizer.Add(self.txt_Password,0, wx.RIGHT, 50)

    	self.btn_Process = wx.Button(self.panel, -1, "&Entrar")
    	self.panel.Bind(wx.EVT_BUTTON,self.OnLogin, self.btn_Process)
    	sizer.Add(self.btn_Process,0, wx.LEFT, 50)

    	self.lbl_label = wx.StaticText(self.panel,-1,"",pos=(180,140))

    	self.panel.SetSizer(sizer)

    def CreateBox (self):
    	self.l2_login = wx.StaticText(self.panel, -1, "CREATE ACCOUNT",pos=(50,220))
    	self.l2_login.SetFont(wx.Font(16, wx.SWISS, wx.NORMAL, wx.BOLD))

    	self.txt2_Username = wx.TextCtrl(self.panel, 1, size=(150, -1),pos=(178,260))
    	self.lbl2_Username = wx.StaticText(self.panel, -1, "Utilizador:",pos=(50,260))

    	self.txt2_Password = wx.TextCtrl(self.panel, 1, size=(150, -1),pos=(178,300),style=wx.TE_PASSWORD)
    	self.lbl2_Password = wx.StaticText(self.panel, -1, "Password:",pos=(50,300))

    	self.lbl2_label = wx.StaticText(self.panel,-1,"",pos=(180,345))

    	self.btn2_Process = wx.Button(self.panel, -1, "&Criar",pos=(50,340))
    	self.panel.Bind(wx.EVT_BUTTON,self.OnCreate, self.btn2_Process)

        self.pic= wx.StaticBitmap(self.panel,pos=(420,80))
        self.pic.SetBitmap(wx.Bitmap("LOGO_ESTIG.png"))

    def Main(self):

        ###USER
        if self.user != None:
            self.label_user = wx.StaticText(self.panel, -1,self.user,pos=(300,2))
            self.label_user.SetFont(wx.Font(8, wx.SWISS, wx.NORMAL, wx.BOLD))
            self.label_user.SetForegroundColour(wx.BLUE)

        ###PORT SCAN
        self.label_init_portscan = wx.StaticText(self.panel, -1, "PORTSCAN",pos=(20,20))
        self.label_init_portscan.SetFont(wx.Font(12, wx.SWISS, wx.NORMAL, wx.BOLD))

        self.input_portscan = wx.TextCtrl(self.panel, 1, size=(150, -1),pos=(20,70))
        self.label_portscan = wx.StaticText(self.panel, -1, "IP Address",pos=(20,50))

        self.show_openports = wx.TextCtrl(self.panel,pos=(190,72),size=(160,110), style=wx.TE_MULTILINE|wx.TE_READONLY)
        self.label_portscan = wx.StaticText(self.panel, -1, "Portos Abertos",pos=(190,50))

        self.portscan_submit= wx.Button(self.panel, -1, "&Scan",pos=(20,110),size=(60,50))
        self.panel.Bind(wx.EVT_BUTTON,self.OnPortScan, self.portscan_submit)

        self.portscan_save= wx.Button(self.panel, -1, "&Save",pos=(110,110),size=(60,50))
        self.panel.Bind(wx.EVT_BUTTON,self.OnPortScanSave, self.portscan_save)

        ###CONEXÕES ACTIVAS
        self.label_init_netstat = wx.StaticText(self.panel, -1, "CONEXÕES ACTIVAS",pos=(400,20))
        self.label_init_netstat.SetFont(wx.Font(12, wx.SWISS, wx.NORMAL, wx.BOLD))

        self.list2=wx.ListCtrl(self.panel,size=(405,155),pos=(400,53),style=wx.LC_REPORT|wx.BORDER_SUNKEN)
        self.list2.Show(True)
        
        self.list2.InsertColumn(0,"STATUS")
        self.list2.InsertColumn(1,"SRC")
        self.list2.InsertColumn(2,"DST")
        self.list2.InsertColumn(3,"PROTO")

        self.list2.SetColumnWidth(0,100)
        self.list2.SetColumnWidth(1,120)
        self.list2.SetColumnWidth(2,120)
        self.list2.SetColumnWidth(3,60)

        self.netstat_submit= wx.Button(self.panel, -1, "&Check",pos=(720,17))
        self.panel.Bind(wx.EVT_BUTTON,self.OnNetStat, self.netstat_submit)

        self.netstat_extract= wx.Button(self.panel, -1, "&Extrair",pos=(620,17))
        self.panel.Bind(wx.EVT_BUTTON,self.saveFileNetStat, self.netstat_extract)
        

        ###ANALISE DE  LOGS
        self.label_init_logs = wx.StaticText(self.panel, -1, "ANÁLISE DE LOGS",pos=(20,210))
        self.label_init_logs.SetFont(wx.Font(12, wx.SWISS, wx.NORMAL, wx.BOLD))

        self.list=wx.ListCtrl(self.panel,size=(500,200),pos=(20,240),style=wx.LC_REPORT|wx.BORDER_SUNKEN)
        self.list.Show(True)

        self.list.InsertColumn(0,"DATE")
        self.list.InsertColumn(1,"TYPE")
        self.list.InsertColumn(2,"DST")
        self.list.InsertColumn(3,"SRC")
        self.list.InsertColumn(4,"FROM")

        self.list.SetColumnWidth(0,120)
        self.list.SetColumnWidth(1,70)
        self.list.SetColumnWidth(2,120)
        self.list.SetColumnWidth(3,105)

        self.logs_path_entry = wx.TextCtrl(self.panel, 1, size=(150, -1),pos=(560,260))
        self.logs_lbl = wx.StaticText(self.panel, -1, "Ficheiro",pos=(560,240))

        self.logs_search= wx.Button(self.panel, -1, "&Procurar",pos=(720,258))
        self.panel.Bind(wx.EVT_BUTTON,self.loadFile, self.logs_search)

        self.logs_submit= wx.Button(self.panel, -1, "&Analisar",pos=(720,359),size=(88,80))
        self.panel.Bind(wx.EVT_BUTTON,self.OnLogs, self.logs_submit)

        self.logs_extract= wx.Button(self.panel, -1, "&Extrair",pos=(563,359),size=(88,80))
        self.panel.Bind(wx.EVT_BUTTON,self.saveFileLogs, self.logs_extract)

        self.logs_graph= wx.Button(self.panel, -1, "&Graph",pos=(563,300),size=(88,50))
        self.panel.Bind(wx.EVT_BUTTON,self.graph, self.logs_graph)

     #Funcão que adiciona elementos a lista
    def addRow2(self,index,status,src,dst,proto):
        pos = self.list2.InsertStringItem(index,status)
        self.list2.SetStringItem(pos,1,src)
        self.list2.SetStringItem(pos,2,dst)
        self.list2.SetStringItem(pos,3,proto)
        if index % 2 == 0:
            self.list2.SetItemBackgroundColour(pos, wx.Colour(200,255,180)) 

    def addRow(self,index,date,type_,src,dst,frome):
        pos = self.list.InsertStringItem(index,date)
        self.list.SetStringItem(pos,1,type_)
        self.list.SetStringItem(pos,2,src)
        self.list.SetStringItem(pos,3,dst)
        self.list.SetStringItem(pos,4,frome)
        if index % 2 == 0:
            self.list.SetItemBackgroundColour(pos, wx.Colour(200,180,255)) 

     #FUNCÃO responsavel pela renderização do gráfico
    def graph(self, event):
        count = self.list.GetItemCount()
        listofsrc = []
        if count > 0:
            for row in range(count):
                if self.list.GetItem(itemId=row, col=4).GetText().strip() != "":
                    listofsrc.append(self.list.GetItem(itemId=row, col=4).GetText())
            if len(listofsrc) > 0:
                labels = []
                sizes = []
                for l in listofsrc:
                    counter = 0
                    if len(labels) > 0:
                        if l in labels:
                            continue
                    for x in listofsrc:
                        if x == l:
                            counter+=1
                    labels.append(l)
                    sizes.append(counter)
                
                fig = plt.gcf()
                fig.canvas.set_window_title('Attacks Graph')
                plt.pie(sizes, labels=labels,autopct='%1.1f%%', shadow=True, startangle=90)
                plt.axis('equal')
                plt.show()

####  --------- EVENTS  ---------- #########
    def OnPortScanSave(self,event):
        if self.show_openports.GetValue().strip() != "" and self.input_portscan.GetValue().strip() != "":
            self.db.connect()
            for port in self.show_openports.GetValue().strip().split(','):
                self.db.addOpenPorts(self.input_portscan.GetValue(),port,self.user)
            #print self.db.getOpenPortsByHost(self.input_portscan.GetValue(),self.user)
    
    def OnClose(self, event):
        dlg = wx.MessageDialog(self,
            "Tem a certeza que pretende sair?",
            "Confirmação", wx.OK|wx.CANCEL|wx.ICON_QUESTION)
        result = dlg.ShowModal()
        dlg.Destroy()
        if result == wx.ID_OK:
            self.Destroy()

    def OnLogin(self, event):
    	username = self.txt_Username.GetValue()
    	password = self.txt_Password.GetValue()
    	if username.strip() == "" or password.strip() == "":
    		self.lbl_label.SetLabel("!!! Campos vazios")
    		self.lbl_label.SetForegroundColour((0,0,255))
    	else:
    		self.db.connect()
    		if self.db.loginAccount(username,password):
    			self.user = username
    			self.hideLoginBox()
    			self.hideCreateBox()
    			self.Main()
    		else:
    			self.lbl_label.SetLabel("!!! Utilizador/Password errados")
    			self.lbl_label.SetForegroundColour((255,0,0))

    def OnCreate(self,event):
    	username = self.txt2_Username.GetValue()
    	password = self.txt2_Password.GetValue()
    	if username.strip() == "" or password.strip() == "":
    		self.lbl2_label.SetLabel("!!! Campos vazios")
    		self.lbl2_label.SetForegroundColour((0,0,255))
    	else:
    		self.db.connect()
    		if self.db.createAccount(username,password):
    			self.lbl2_label.SetLabel("Novo Utilizador Criado")
    			self.lbl2_label.SetForegroundColour((0,200,0))
    		else:
    			self.lbl2_label.SetLabel("!!! Utilizador já existe")
    			self.lbl2_label.SetForegroundColour((255,0,0))


   	"""Eventos de inicialização das Threads"""
    def OnPortScan(self, event):
        if not self.portscan and self.input_portscan.GetValue().strip() != "":
        	self.portscan = PortScanThread(self,self.input_portscan.GetValue(),1000)
    
    def OnNetStat(self,event):
        if not self.netstat:
            self.netstat = NetStatThread(self)
    def OnLogs(self,event):
    	if not self.logs and self.logs_path_entry.GetValue().strip() != "":
            self.logs = LogsThread(self,self.logs_path_entry.GetValue())

    """Evento responsavel por receber os sinais das threads"""
    def OnResult(self, event):
        """Show Result status."""
        if event.data['type'] == 4:
            if len(event.data['connections']) > 0:
                i = 0
                self.list2.DeleteAllItems()
                for conn in event.data['connections']:
                    self.addRow2(i,conn['STATUS'],conn['SRC']+":"+conn['SRC_PORT'],conn['DEST']+":"+conn['DEST_PORT'],conn['PROTO'])
                    i+=1
            self.netstat = None
        elif event.data['type'] == 5:
            if len(event.data['logs']) > 0:
                i = 0
                self.list.DeleteAllItems()
                for log in event.data['logs']:
                    self.addRow(i,log['DATE'],log['TYPE'].replace("UFW","").strip(),log['DST'],log['SRC'],self.gi.country_name_by_addr(log['SRC']))
                    i+=1
            self.logs = None
        elif event.data['type'] == 3:
            if len(event.data['ports']):
                output = ""
                i = 0
                for port in event.data['ports']:
                    if i == 0:
                        output+=str(port)
                    else:
                        output+=","+str(port)
                    i+=1
                self.show_openports.SetValue(output)
            self.portscan = None


    def saveFileLogs(self, event=None):

        count = self.list.GetItemCount()
        if count > 0:
            saveFileDialog = wx.FileDialog(self,"Save As", "", "","CSV files (*.csv)|*.csv",wx.FD_SAVE | wx.FD_OVERWRITE_PROMPT)
            saveFileDialog.ShowModal()

            file = open(saveFileDialog.GetPath()+".csv", "w+")
            count = self.list.GetItemCount()
            i = 0
            for row in range(count):
                if i == 0:
                    file.write('DATE,TYPE,DST,SRC,FROM')
                file.write("\n"+self.list.GetItem(itemId=row, col=0).GetText()+","+self.list.GetItem(itemId=row, col=1).GetText()+","+self.list.GetItem(itemId=row, col=2).GetText()+","+self.list.GetItem(itemId=row, col=3).GetText()+","+self.list.GetItem(itemId=row, col=4).GetText())
                i+=1
            file.close()
            saveFileDialog.Destroy()

    def saveFileNetStat(self, event=None):

        count = self.list2.GetItemCount()
        if count > 0:
            saveFileDialog = wx.FileDialog(self,"Save As", "", "","CSV files (*.csv)|*.csv",wx.FD_SAVE | wx.FD_OVERWRITE_PROMPT)
            saveFileDialog.ShowModal()

            file = open(saveFileDialog.GetPath()+".csv", "w+")
            count = self.list2.GetItemCount()
            i = 0
            for row in range(count):
                if i == 0:
                    file.write('STATUS,SRC,DST,PROTO')
                file.write("\n"+self.list2.GetItem(itemId=row, col=0).GetText()+","+self.list2.GetItem(itemId=row, col=1).GetText()+","+self.list2.GetItem(itemId=row, col=2).GetText()+","+self.list2.GetItem(itemId=row, col=3).GetText())
                i+=1
            file.close()
            saveFileDialog.Destroy()

    def loadFile(self,event=None,FILES="Log files (*.log)|*.log"):
    	openFileDialog = wx.FileDialog(self, "Open", "", "",FILES, wx.FD_OPEN | wx.FD_FILE_MUST_EXIST)
    	openFileDialog.ShowModal()
        self.logs_path_entry.SetValue(openFileDialog.GetPath())
    	openFileDialog.Destroy()

def main():
	app = wx.App(redirect=True)
	top = Frame()
	top.Show()
	app.MainLoop()

if __name__ == "__main__":
	main()